{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbcaf554-dd74-41d4-9a01-6fdb6021d24a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "model_name = './gpt2'\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9346d98d-be4c-4b2b-9330-b68392b83b62",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "dataset_name = './sst2'\n",
    "ds = load_dataset(dataset_name)\n",
    "ds_train, ds_val = ds['train'], ds['validation']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "602f4b52-a134-4a5a-a284-6984b0cdf1a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4174501-5ccb-4ebe-bb86-b344146f4138",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cab7d134-d766-4513-a6b1-6f4e5177ec67",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_train[6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55cd60bc-43b5-41d3-965a-78230971b334",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_train[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94a8cd2b-ad43-49b0-854b-f70ad899d8fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(batch):\n",
    "    return tokenizer(batch['sentence'])\n",
    "\n",
    "map_kwargs = {\n",
    "    'batched': True,\n",
    "    'batch_size': 512,\n",
    "    'remove_columns': ['idx', 'sentence', 'label']\n",
    "}\n",
    "\n",
    "tokenized_dataset_train = ds_train.map(tokenize, **map_kwargs)\n",
    "tokenized_dataset_val = ds_val.map(tokenize, **map_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7cc0c03-04f4-4c52-9838-b9cd4fb002b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_dataset_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bb1c731-2a87-4a64-a674-aff7caf7c137",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_dataset_train[5:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b467dbff-e9c8-49f4-b6c6-63beaa427d36",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, seq in enumerate(tokenized_dataset_train[5:10]['input_ids']):\n",
    "    print(f'{i+1}: {tokenizer.decode(seq)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d64e5141-5baf-4dd4-baf4-c7a4052e63bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(tokenized_dataset_train), len(tokenized_dataset_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d3d7481-8d23-4a78-b574-c8065b517930",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_dataset_train = tokenized_dataset_train.filter(lambda x: len(x['input_ids']) > 5)\n",
    "tokenized_dataset_val = tokenized_dataset_val.filter(lambda x: len(x['input_ids']) > 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a3cf4c6-cfad-46a7-8fd1-e2e0f2b1ea95",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(tokenized_dataset_train), len(tokenized_dataset_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "490c0186-8cf4-4950-986c-d57fc51b5ba3",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_dataset_train.set_format(type='torch')\n",
    "tokenized_dataset_val.set_format(type='torch')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d86126ef-ab5e-4fda-a9ba-bd348ef15d8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_dataset_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c91f606d-71b5-4bc2-a343-8ed4b52657ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_dataset_train[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a49499c-4721-4353-91b7-a5d25b1d2515",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.pad_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "458aef5f-fc27-455d-ab75-884e2d79eadb",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be3517d1-2e0b-4987-b15d-ec4769aa98fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f354dd1e-2342-4d79-8ea9-f0b46d5788d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from transformers import DataCollatorForLanguageModeling\n",
    "data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False) # labels\n",
    "\n",
    "dataloader_params = {\n",
    "    'batch_size': 16,\n",
    "    'collate_fn': data_collator\n",
    "}\n",
    "\n",
    "train_dataloader = DataLoader(tokenized_dataset_train, **dataloader_params)\n",
    "val_dataloader = DataLoader(tokenized_dataset_val, **dataloader_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe54ccfd-8d6f-4697-8f3b-4d56a7836db4",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(train_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "641edbd5-a5db-420f-ad19-2aa1b6af9152",
   "metadata": {},
   "outputs": [],
   "source": [
    "3088 * 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0d0515b-eb79-47d6-9a3e-aa954ae5e7fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = next(iter(train_dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf60bd51-3562-4815-844a-43d360c275b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78583cc8-a930-4097-884b-5cef3bc69343",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch['input_ids'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af63666a-50c7-4d73-a9b8-711e3e485340",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch['input_ids'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5cb28f2-8b36-4383-9346-8fa1d0f636d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch['labels'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7fa8819-ec8d-4e21-9b03-09b8e50f5bd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch['attention_mask'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dcb30f2-eff6-418f-a12d-f5916b8b2f5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)\n",
    "num_epochs = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37f1918c-1d79-49c7-a5fa-47660fa76232",
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate(epoch):\n",
    "    model.eval()\n",
    "    total_loss = 0.0\n",
    "    for i, batch in enumerate(val_dataloader):\n",
    "        batch = batch.to(device)\n",
    "        with torch.no_grad():\n",
    "            outputs = model(**batch)\n",
    "            loss = outputs.loss # 损失\n",
    "            total_loss += loss.item()\n",
    "    print(f'val_loss at {epoch} epoch:', total_loss / len(val_dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7355010e-4db4-4c2e-b3ef-aeda47a13306",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model.to(device)\n",
    "validate(0)\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    for i, batch in enumerate(train_dataloader):\n",
    "        batch = batch.to(device)\n",
    "        outputs = model(**batch)\n",
    "        loss = outputs.loss\n",
    "        print(f'Loss: {loss.item()}')\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    validate(epoch+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e28250b-fc90-490c-b067-02f92060a3c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained('./gpt2-sft')\n",
    "tokenizer.save_pretrained('./gpt2-sft')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "444b1810-7469-40b4-ba50-7c70b06482e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline, set_seed\n",
    "from pprint import pprint\n",
    "g = pipeline('text-generation', model='./gpt2-sft')\n",
    "set_seed(1337)\n",
    "pprint(g(\"Hi, this is all terribly\", max_length=30, num_return_sequences=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a15981e-1fcf-4f18-bb8f-4ce7d183ed9d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
